cd ~/Dara/yolov9-main

# Create script to find layer indices
cat > find_layers.py << 'EOF'
import torch
from models.yolo import DetectionModel

print("Finding P3, P4, P5 layer indices for YOLOv9...")
print("="*70)

try:
    model = DetectionModel('models/yolov9-c.yaml')
    dummy_input = torch.randn(1, 3, 640, 640)
    
    y = []
    candidates = {'P3': [], 'P4': [], 'P5': []}
    
    for i, m in enumerate(model.model):
        if m.f != -1:
            dummy_input = y[m.f] if isinstance(m.f, int) else \
                [dummy_input if j == -1 else y[j] for j in m.f]
        dummy_input = m(dummy_input)
        y.append(dummy_input if m.i in model.save else None)
        
        if isinstance(dummy_input, torch.Tensor) and len(dummy_input.shape) == 4:
            H, W = dummy_input.shape[2:]
            C = dummy_input.shape[1]
            
            if H == 80 and W == 80:
                candidates['P3'].append((m.i, C))
                print(f"Layer {m.i:3d}: shape={dummy_input.shape} | P3 CANDIDATE (80×80)")
            elif H == 40 and W == 40:
                candidates['P4'].append((m.i, C))
                print(f"Layer {m.i:3d}: shape={dummy_input.shape} | P4 CANDIDATE (40×40)")
            elif H == 20 and W == 20:
                candidates['P5'].append((m.i, C))
                print(f"Layer {m.i:3d}: shape={dummy_input.shape} | P5 CANDIDATE (20×20)")
    
    print("="*70)
    print("\n📋 SUMMARY:")
    print("-"*70)
    
    if candidates['P3']:
        idx, ch = candidates['P3'][-1]  # Take last one
        print(f"P3: Layer {idx:3d} with {ch:4d} channels")
    
    if candidates['P4']:
        idx, ch = candidates['P4'][-1]
        print(f"P4: Layer {idx:3d} with {ch:4d} channels")
    
    if candidates['P5']:
        idx, ch = candidates['P5'][-1]
        print(f"P5: Layer {idx:3d} with {ch:4d} channels")
    
    print("-"*70)
    
    # Generate code
    print("\n📝 UPDATE models/yolo.py with these indices:")
    print("-"*70)
    if candidates['P3'] and candidates['P4'] and candidates['P5']:
        p3_idx = candidates['P3'][-1][0]
        p4_idx = candidates['P4'][-1][0]
        p5_idx = candidates['P5'][-1][0]
        
        p3_ch = candidates['P3'][-1][1]
        p4_ch = candidates['P4'][-1][1]
        p5_ch = candidates['P5'][-1][1]
        
        print(f"""
# In extract_features_for_stata():
if m.i == {p3_idx}:  # P3
    features.append(x)
elif m.i == {p4_idx}:  # P4
    features.append(x)
elif m.i == {p5_idx}:  # P5
    features.append(x)

# In setup_stata():
yolo_channels = [{p3_ch}, {p4_ch}, {p5_ch}]  # P3, P4, P5
""")
    else:
        print("⚠️  Could not find all scales. Check your model configuration.")
    
    print("-"*70)

except Exception as e:
    print(f"\n✗ Error: {e}")
    print("\nMake sure:")
    print("  1. models/yolov9-c.yaml exists")
    print("  2. YOLOv9 model is properly installed")
EOF

python find_layers.py